from fastapi.exceptions import HTTPException
from contextlib import contextmanager
import time
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint


@contextmanager
def HttpResponse():
    exception = None
    try:
        yield
    except HTTPException as e:
        exception = HTTPException(status_code=e.status_code, detail=e.detail, headers=e.headers)
    except Exception as e:
        exception = HTTPException(status_code=500, detail=str(e))
    finally:
        if exception:
            raise exception


class CustomMiddleware(BaseHTTPMiddleware):

    async def dispatch(self, request, call_next):
        try:
            start_time = time.time()
            responser = await call_next(request)
            process_time = round(time.time() - start_time, 4)
            responser.headers["X-Process-Time"] = f"{process_time} (s)"
            return responser
        except Exception as e:
            raise HTTPException(status_code=500,
                                detail={
                                    "message": f"{request.url.path} 接口执行异常, 联系工程师修复.",
                                    "trace": str(e)})
